import requests
import re
import os
import time
import win32api
import random

def crawler():#爬虫
    # 定义请求头
    headers = {"user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                             "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/135.0.0.0 Safari/537.36 Edg/135.0.0.0"}

    # 遍历所有url
    urls = ['https://www.gushici.net/index.html']  # 网站第一页
    for i in range(2, 101, 1):
        url = 'https://www.gushici.net/index_{}.html'.format(i)
        urls.append(url)

    # 存储路径
    output_dir = r'F:\PoemsGet'
    os.makedirs(output_dir, exist_ok=True)

    # PageCount计算已爬取的页数,allPoem计算已爬取数量，errCount计算问题数量
    PageCount = 0
    allPoem = 0
    errCount = 0

    for url in urls:
        response = requests.get(url, headers=headers)
        content = response.content.decode('utf-8')

        # 在html上提取数据
        titles = re.findall('<b>(.*?)</b>', content, re.DOTALL)
        dynasties = re.findall('<p class="source"><a.*?>(.*?)</a>', content, re.DOTALL)
        authors = re.findall('<p class="source"><a.*?<a.*?>(.*?)</a>', content, re.DOTALL)
        poems = re.findall('<div class="gushici-box-text">(.*?)</div>', content, re.DOTALL)
        time.sleep(3)

        # 处理诗句内容
        new_poems = []
        for poem in poems:
            new_poem = re.sub('<a.*?>', "", poem)
            new_poem = re.sub('</a>&nbsp;', "", new_poem)
            new_poem = re.sub('<br />', "", new_poem)
            new_poem = new_poem.strip()
            new_poems.append(new_poem)

        for title, dynasty, author, new_poem in zip(titles, dynasties, authors, new_poems):
            try:
                gushici = {
                    "title": title,
                    "dynasty": dynasty,
                    "author": author,
                    "poem": new_poem
                }
                file_name = "{}.txt".format(gushici['title'])
                file_path = os.path.join(output_dir, file_name)
                print(gushici['title'])
                with open(file_path, 'w', encoding='utf-8') as f:
                    for value in gushici['title'], gushici['dynasty'], gushici['author'], gushici['poem']:
                        f.write(value)
                        f.write('.')
                        f.write('\n')
                del gushici
                allPoem += 1
            except:
                errCount += 1
                continue
        PageCount += 1
        print("————已经爬取第{}页————".format(PageCount))

    print("————总爬取数量：{}\n".format(allPoem))
    print("————因错误未保存的古诗词数量：{}\n".format(errCount))

def judgement(): #判断是否已爬取
    if len(os.listdir('F:/PoemsGet')) == 0:
        print('开始爬取古诗词，请稍等…\n')
        crawler() #爬虫
    else:
        print('已有古诗词，请移步AI语音功能')

def open_sovits(): #启动AI语音
    win32api.ShellExecute(0,'open',
                          'F:/GPT-SoVITS-v2/GPT-SoVITS-v2-240821/v2-webui.bat',"","",1)

def select_poem():
    if len(os.listdir('F:/PoemsGet')) == 0:
        print('无古诗词，请先按1获取！\n')
    else:
        source_folder = r'F:/PoemsGet' #存古诗的文件夹
        #获取所有已保存的古诗
        allfiles = [f for f in os.listdir(source_folder) if
                    f.endswith('.txt') and os.path.isfile(os.path.join(source_folder, f))]

        #随机抽取
        selected_file = random.choice(allfiles)
        file_path = os.path.join(source_folder, selected_file)

        #打印抽选内容
        with open(file_path, "r", encoding="utf-8") as f:
            content = f.read()
            print(content)
        del allfiles
        del selected_file
        del source_folder

def main_loop():
    while True:
        print("=" * 30)
        print("请选择要执行的功能：")
        print("1. 获取古诗词")
        print("2. 启动AI语音")
        print("3. 随机选取古诗词")
        print("0. 退出程序")
        print("=" * 30)
        choice = input("请输入选项：").strip()
        print('\n')

        if choice == '1':
            judgement()
        elif choice == '2':
            open_sovits()
        elif choice == '3':
            select_poem()
        elif choice == '0':
            print("程序已退出！")
            break
        else:
            print("无效输入，请重新输入！")

if __name__ == '__main__':
    print('请选择要执行的功能...')
    main_loop()